import os
import cv2
import numpy as np
import tqdm

path = '/media/hjh/workdir/0_Deep_Learning/0_top1/DOTA/labelTxt'
save_path = '/media/hjh/workdir/0_Deep_Learning/0_top1/DOTA/labels'
write_file = open('train.txt', 'w')

width, height = 600, 600
category = [
        'plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle',
        'ship', 'tennis-court', 'basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout',
        'harbor', 'swimming-pool', 'helicopter'
        ]

list_dir = os.listdir(path)
for i in tqdm.tqdm(list_dir):
    txt_file = open(os.path.join(path, i), 'r')
    lines = txt_file.readlines()
    #
    valid_pts = []
    for line in lines:
        obj = line.split(' ') # list object
        if len(obj)>8:
            x1 = min(max(float(obj[0]), 0), width-1)
            y1 = min(max(float(obj[1]), 0), height-1)
            x2 = min(max(float(obj[2]), 0), width-1)
            y2 = min(max(float(obj[3]), 0), height-1)
            x3 = min(max(float(obj[4]), 0), width-1)
            y3 = min(max(float(obj[5]), 0), height-1)
            x4 = min(max(float(obj[6]), 0), width-1)
            y4 = min(max(float(obj[7]), 0), height-1)
            # TODO: filter small instances
            xmin = max(min(x1, x2, x3, x4), 0)
            xmax = max(x1, x2, x3, x4)
            ymin = max(min(y1, y2, y3, y4), 0)
            ymax = max(y1, y2, y3, y4)
            box_w = xmax - xmin
            box_h = ymax - ymin
            if (box_w > 10) and (box_h > 10):
                box = [x1, y1, x2, y2, x3, y3, x4, y4, obj[8]]
                valid_pts.append(box)
    if len(valid_pts):
        new_path = os.path.join(save_path, i)
        new_file = open(new_path, 'w')
        for p in valid_pts:
            p = str(p)
            write_str = p + '\n'
            new_file.write(write_str)
        new_file.close()
        # print(new_path, '\t写入成功')
                
    else:
        # print(i, '\t没有object')
        pass